package xyz.xuetu.www.rpc.registry;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.cron.CronUtil;
import cn.hutool.cron.task.Task;
import cn.hutool.json.JSONUtil;
import io.etcd.jetcd.*;
import io.etcd.jetcd.options.GetOption;
import io.etcd.jetcd.options.PutOption;
import io.etcd.jetcd.watch.WatchEvent;
import io.vertx.core.impl.ConcurrentHashSet;
import lombok.extern.slf4j.Slf4j;
import xyz.xuetu.www.rpc.config.RegistryConfig;
import xyz.xuetu.www.rpc.model.ServiceMetaInfo;

import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ExecutionException;
import java.util.stream.Collectors;

@Slf4j
public class EtcdRegistry implements Registry {

    private Client client;

    private KV kvClient;

    /**
     * 根节点
     */
    private static final String ETCD_ROOT_PATH = "/rpc/";

    /**
     * 本机注册的节点 key 集合(用于维护续期)
     */
    private final Set<String> localRegisterZNodeKeySet = new HashSet<>();

    /**
     * 注册中心服务缓存
     */
    private final RegistryServiceCache registryServiceCache = new RegistryServiceCache();

    /**
     * 正在监听的 key 集合
     */
    private final Set<String> watchingKeySet = new ConcurrentHashSet<>();

    @Override
    public void init(RegistryConfig registryConfig) {
         client = Client.builder()
                .endpoints(registryConfig.getAddress())
                .connectTimeout(Duration.ofMillis(registryConfig.getTimeout()))
                .build();
         kvClient = client.getKVClient();
         // 开启心跳检测和续签机制
         heartBeat();
    }

    @Override
    public void register(ServiceMetaInfo serviceMetaInfo) throws ExecutionException, InterruptedException {
        // 创建 Lease 和 KV 客户端
        Lease leaseClient = client.getLeaseClient();
        // 创建一个 30 秒的租约
        long leaseId = leaseClient.grant(30).get().getID();
        // 设置要存储的键值对
        String registerKey = ETCD_ROOT_PATH + serviceMetaInfo.getServiceNodeKey();
        ByteSequence byteSequenceKey = ByteSequence.from(registerKey, StandardCharsets.UTF_8);
        ByteSequence byteSequenceValue = ByteSequence.from(JSONUtil.toJsonStr(serviceMetaInfo), StandardCharsets.UTF_8);
        // 将键值对于租约关联起来，并设置过期时间
        PutOption putOption = PutOption.builder().withLeaseId(leaseId).build();
        kvClient.put(byteSequenceKey, byteSequenceValue, putOption).get();
        // 添加节点信息到本地缓存
        localRegisterZNodeKeySet.add(registerKey);
    }

    @Override
    public void unRegister(ServiceMetaInfo serviceMetaInfo) {
        String registerKey = ETCD_ROOT_PATH + serviceMetaInfo.getServiceNodeKey();
        ByteSequence byteSequenceKey = ByteSequence.from(registerKey, StandardCharsets.UTF_8);
        kvClient.delete(byteSequenceKey);
        // 从本地缓存移除
        localRegisterZNodeKeySet.remove(registerKey);
    }

    @Override
    public List<ServiceMetaInfo> serviceDiscovery(String serviceKey) {
        // 有限从缓存获取服务
        List<ServiceMetaInfo> cacheServiceMetaInfoList = registryServiceCache.readCache();
        if(cacheServiceMetaInfoList != null) {
            return cacheServiceMetaInfoList;
        }

        // 前缀搜索
        String searchPrefix = ETCD_ROOT_PATH + serviceKey + "/";
        try {
            // 前缀查询
            GetOption getOption = GetOption.builder().isPrefix(true).build();
            ByteSequence byteSequenceSearchPrefix = ByteSequence.from(searchPrefix, StandardCharsets.UTF_8);
            List<KeyValue> keyValueList = kvClient.get(byteSequenceSearchPrefix, getOption).get().getKvs();
            // 解析服务信息
            List<ServiceMetaInfo> serviceMetaInfoList = keyValueList.stream()
                    .map(keyValue -> {
                        String key = keyValue.getKey().toString(StandardCharsets.UTF_8);
                        // 监听 key 的变化
                        watch(key);
                        String value = keyValue.getValue().toString(StandardCharsets.UTF_8);
                        return JSONUtil.toBean(value, ServiceMetaInfo.class);
                    })
                    .collect(Collectors.toList());

            // 写入服务缓存
            registryServiceCache.writeCache(serviceMetaInfoList);
            return serviceMetaInfoList;
        } catch (InterruptedException | ExecutionException e) {
            log.error("获取服务列表失败 {}", e.getMessage());
            throw new RuntimeException(String.format("获取服务列表失败 {}", e.getMessage()));
        }
    }

    @Override
    public void destroy() {
        log.info("当前节点下线");
        // 遍历本节点所有的 key
        localRegisterZNodeKeySet.forEach(key -> {
            try {
                ByteSequence byteSequenceKey = ByteSequence.from(key, StandardCharsets.UTF_8);
                kvClient.delete(byteSequenceKey).get();
            } catch (InterruptedException | ExecutionException e) {
                throw new RuntimeException(key + "节点下线失败", e);
            }
        });

        // 释放资源
        if(kvClient != null) {
            kvClient.close();
        }
        if(client == null) {
            client.close();
        }
    }

    @Override
    public void heartBeat() {
        // 10 秒续签一次
        CronUtil.schedule("*/10 * * * * *", new Task() {
            @Override
            public void execute() {
                // 遍历本节点所有的 key
                localRegisterZNodeKeySet.forEach(key -> {
                    try {
                        ByteSequence byteSequenceKey = ByteSequence.from(key, StandardCharsets.UTF_8);
                        List<KeyValue> keyValueList = kvClient.get(byteSequenceKey)
                                .get()
                                .getKvs();
                        // 该节点未过期，重新注册(相当于续签)
                        if(CollUtil.isNotEmpty(keyValueList)) {
                            KeyValue keyValue = keyValueList.get(0);
                            String value = keyValue.getValue().toString(StandardCharsets.UTF_8);
                            ServiceMetaInfo serviceMetaInfo = JSONUtil.toBean(value, ServiceMetaInfo.class);
                            register(serviceMetaInfo);
                        }
                    } catch (InterruptedException | ExecutionException e) {
                        throw new RuntimeException(key + "续签失败", e);
                    }
                });
            }
        });
        // 支持秒级别定时任务
        CronUtil.setMatchSecond(true);
        CronUtil.start();
    }

    @Override
    public void watch(String serviceNodeKey) {
        Watch watchClient = client.getWatchClient();
        // 之前未被监视的，开始监听
        boolean isAdd = watchingKeySet.add(serviceNodeKey);
        if(isAdd) {
            ByteSequence byteSequenceKey = ByteSequence.from(serviceNodeKey, StandardCharsets.UTF_8);
            watchClient.watch(byteSequenceKey, response -> {
                for (WatchEvent event : response.getEvents()) {
                    switch (event.getEventType()) {
                        // key删除时触发
                        case DELETE:
                            // 清除注册服务缓存
                            registryServiceCache.clearCache();
                            break;
                        case PUT:
                        default:
                            break;
                    }
                }
            });
        }
    }
}
